# -*- coding: utf-8 -*-
"""
Created on Wed Apr 10 09:57:49 2019

@author: Fsl
"""

import torch
from torchvision import models
import torch.nn as nn
from model.resnet import resnet101, resnet34, resnet50
from torch.nn import functional as F
import torchsummary
from torch.nn import init
from model.Aspp import build_aspp

BatchNorm = nn.BatchNorm2d
up_kwargs = {'mode': 'bilinear', 'align_corners': True}


# **********边界图辅助分割，两个deeplabV3+的分支，引入的deeplab的resnet模块

class CPFNet(nn.Module):
    #  'BaseNet': CPFNet(out_planes=args.num_classes)
    def __init__(self, out_planes=1, ccm=True, norm_layer=nn.BatchNorm2d, is_training=True, expansion=2, base_channel=32):
        super(CPFNet, self).__init__()
        self.out_chanel = out_planes
        self.in_de_c = 128  # 特征解码器对编码器的特征统一降通道
        self.network = "resnet34"

        if self.network == "resnet101":
            self.backbone = resnet101(pretrained=True)
            self.in_c = [64, 256, 512, 1024, 2048]
        
        if self.network == "resnet34":
            self.backbone = resnet34(pretrained=True)
            self.in_c = [64, 64, 128, 256, 512]

        # aspp的通道
        # Decoder的aspp输出的通道
        self.aspp_s = build_aspp(self.in_c[-1], 256)
        self.aspp_e = build_aspp(self.in_c[-1], 256)
        
        self.edge_4 = self.backbone.layer4
        self.decoder_e = Decoder_V3()
        self.decoder_s = Decoder_V3()
        
        # 边界注意力模块
        self.BaSm2 = BaSm(64)
        self.BaSm5 = BaSm(256)
        
        self.sigmoid_e = torch.nn.Sigmoid()
        self.sigmoid_s = torch.nn.Sigmoid()
        
    def forward(self, x):
        h, w = x.size()[2:]
        x = self.backbone.conv1(x)
        x = self.backbone.bn1(x)
        x1 = self.backbone.relu(x)      # 1/2  64
        x = self.backbone.maxpool(x1)
        #                                model     resnet34       resnet101
        x2 = self.backbone.layer1(x)   # 1/4      64              256
        x3 = self.backbone.layer2(x2)  # 1/8      128             512
        x4 = self.backbone.layer3(x3)  # 1/8      256             1024
        x5 = self.backbone.layer4(x4)  # 1/8      512             2048
        
        # print("x2.shape :", x2.shape)
        # print("x3.shape :", x3.shape)
        # print("x4.shape :", x4.shape)
        # print("x5.shape :", x5.shape)
        
        s5 = self.aspp_s(x5)  # 分割
        e5 = self.aspp_e(x5)  # 边界
        
        # 获取边界预测图
        edge = self.decoder_e(e5, x2)
        edge = F.interpolate(edge, size=[h, w], mode='bilinear', align_corners=True)
        out_edge = self.sigmoid_e(edge)
        
        # print("edge.shape", edge.shape)
        # print("out_edge.shape", out_edge.shape)
        
        ed2 = self.BaSm2(out_edge, x2)
        ed5 = self.BaSm5(out_edge, s5)
        
        # 获取分割预测图
        seg = self.decoder_s(ed5, ed2)
        seg = F.interpolate(seg, size=[h, w], mode='bilinear', align_corners=True)
        out_seg = self.sigmoid_s(seg)
        
        # print("seg.shape", seg.shape)
        # print("out_seg.shape", out_seg.shape)
        
        return out_seg, out_edge
    
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_uniform_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                init.normal_(m.weight.data, 1.0, 0.02)
                init.constant_(m.bias.data, 0.0)


class Decoder_V3(nn.Module):
    def __init__(self):
        super(Decoder_V3, self).__init__()
        low_level_inplanes = 64
        self.conv1 = nn.Conv2d(low_level_inplanes, 48, 1, bias=False)
        self.bn1 = BatchNorm(48)
        self.relu = nn.ReLU()
        self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
                                       BatchNorm(256),
                                       nn.ReLU(),
                                       nn.Dropout(0.5),
                                       nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
                                       BatchNorm(256),
                                       nn.ReLU(),
                                       nn.Dropout(0.1),
                                       nn.Conv2d(256, 1, kernel_size=1, stride=1))
        self._init_weight()

    def forward(self, x, low_level_feat):
        # print(low_level_feat.size())
        # print(x.size())
        low_level_feat = self.conv1(low_level_feat)
        low_level_feat = self.bn1(low_level_feat)
        low_level_feat = self.relu(low_level_feat)

        x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat((x, low_level_feat), dim=1)
        x = self.last_conv(x)

        return x

    def _init_weight(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.kaiming_normal_(m.weight)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()


class BaSm(nn.Module):
    def __init__(self, inchannels):
        super(BaSm, self).__init__()
        self.sigmoid = nn.Sigmoid()
        self.conv3x3 = nn.Conv2d(inchannels + 1, inchannels, 3, padding=1, bias=False)
        
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_uniform_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                init.normal_(m.weight.data, 1.0, 0.02)
                init.constant_(m.bias.data, 0.0)
    
    def forward(self, edge, x):
        # 边界压缩
        # x Sigmoid自注意力
        # 边界注意力
        # 拼接，融合
        edge = F.interpolate(edge, x.size()[2:], **up_kwargs)
        x1 = x * self.sigmoid(x)
        x2 = edge * x
        fusion = x + x1 + x2
        fusion = torch.cat([edge, fusion], dim=1)
        fusion = self.conv3x3(fusion)
        return fusion


BaseNet_version = "V9"

if __name__ == '__main__':
    print(BaseNet_version)
    print(" Using BaseNet_V9 : 边界图辅助分割，两个deeplabV3+的分支，引入的deeplab的resnet模块")
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    model = CPFNet()
    # print(model)
    # model = torch.nn.DataParallel(model).cuda()
    model.cuda()
    with torch.no_grad():
        model.eval()
        torchsummary.summary(model, (3, 256, 256))
